/**
 * @file Watershed.cpp
 * @author enemy1205 (enemy1205@qq.com)
 * @brief 分水岭算法的原理实现
 * @date 2021-09-08
 */
#include "Watershed.h"

using namespace Segmentation;

/**
 * @brief 使用说明
 * @return
 */
void WaterShed::help() {
    cout<<"该代码为分水岭算法的简单实现:"<<endl;
    cout<<"大致原理上是根据颜色明度所分割,所以请不要刁难它，尽量选一些简单的图进行测试"<<endl;
    cout<<"开始要先输入测试图片路径(不带“”),回车"<<endl;
    cout<<"再要将图片分割成几个区域，即在对应区域用鼠标画上一笔"<<endl;
    cout<<"滑条可控制画笔粗细,按c可清除重画"<<endl;
    cout<<"画完后按esc即可显示分割后图像"<<endl;
}
/**
 * @brief 鼠标回调函数
 * @param event 鼠标事件，自动给定
 * @param x,y 鼠标所处坐标，自动给定
 * @param flag 回调函数格式要求
 * @param param 鼠标响应处理的参数
 * @return
 */
void onMouse(int event, int x, int y, int flag, void *param) {
    pair<Point, bool> &Pt = *static_cast<pair<Point, bool> *>(param);
    switch (event) {
        case EVENT_MOUSEMOVE: {
            Pt.first = Point(x, y);//位置点不断随鼠标刷新
        }
            break;
        case EVENT_LBUTTONDOWN: {
            Pt.second = true;//绘画状态打开
            Pt.first = Point(x, y);
        }
            break;
        case EVENT_LBUTTONUP: {
            Pt.second = false;//绘画状态关闭
        }
            break;
        default:
            break;
    }
}

/**
 * @brief 构造函数
 * @param path 输入待处理的源图路径
 * @return
 */
WaterShed::WaterShed(string &path) {
    this->src = imread(path);
    this->drawImg = src.clone();
    this->markersImg = Mat::zeros(src.rows, src.cols, CV_32S);
    this->resMask = this->markersImg.clone();
    this->result = Mat::zeros(src.rows, src.cols, CV_8UC3);
    this->basins = 0;
    this->thickness = 2;
    this->drawColor = Scalar(125, 10, 241);
    this->help();
}

/**
 * @brief 给定图片进行初始化
 * @param frame 给定图片
 * @return
*/
void WaterShed::Init(const Mat &frame) {
    this->src = frame;
    this->drawImg = src.clone();
    this->markersImg = Mat::zeros(src.rows, src.cols, CV_32S);
    this->resMask = this->markersImg.clone();
    this->result = Mat::zeros(src.rows, src.cols, CV_8UC3);
    this->basins = 0;
    this->thickness = 2;
    this->drawColor = Scalar(125, 10, 241);
}

/**
 * @brief 鼠标绘制种子点
 * @return
 */
void WaterShed::markers() {
    Point CurPoint;
    bool Drawing = false;
    auto param = make_pair(CurPoint, Drawing);
    namedWindow("Mouse Event", 0);
    setMouseCallback("Mouse Event", onMouse, (void *) &param);
    namedWindow("slide bar", 0);
    createTrackbar("thickness", "slide bar", &thickness, 100, nullptr);
    int key;
    bool lastDrawing = false;
    while (true) {
        CurPoint = param.first;
        Drawing = param.second;
        if (Drawing) {
            circle(drawImg, CurPoint, 0, drawColor, thickness + 1);
            circle(markersImg, CurPoint, 0, Scalar(basins + 1), thickness + 1);
        } else {
            //lastDrawing保证鼠标松开只计数一次，否则松开阶段一直计数
            if (lastDrawing) {
                basins++;//连续绘制一次后标志位发生改变
            }
        }
        imshow("Mouse Event", drawImg);
        key = waitKey(1);
        //退出绘图，开始分割
        if (key == 27) {
            break;
        }
        //清除所绘制内容
        if (key == 'c') {
            cout << "clear" << endl;
            drawImg = src.clone();
            markersImg = Mat::zeros(src.rows, src.cols, CV_32S);
            basins = 0; //分割区域标志位归0
        }
        lastDrawing = Drawing;
    }
}

/**
 * @brief 计算图像灰度函数的近似梯度
 * @return Mat 灰度梯度矩阵
 */
Mat WaterShed::preProcess() {
    Mat gray;
    cvtColor(src, gray, COLOR_BGR2GRAY);
    //bilateralFilter(gray,gray,0,10,5);
    GaussianBlur(gray, gray, {3, 3}, 1);
    //计算图像灰度函数的近似梯度
    Mat g_x, g_y;
    Sobel(gray, g_x, CV_64F, 1, 0);//x方向灰度变化矩阵
    Sobel(gray, g_y, CV_64F, 0, 1);//y方向灰度变化矩阵
    Mat grad;//反映整体灰度变化的矩阵
    magnitude(g_x,g_y,grad);//平方和再开方
    double minv = 0.0, maxv = 0.0;
    double *minp = &minv;
    double *maxp = &maxv;
    minMaxIdx(grad, minp, maxp);//找寻梯度矩阵里面的最大最小值
    //范围内归一化
    grad = (grad - minv) * 255 / (maxv - minv);
    return grad;
}

/**
 * @brief 根据种子及灰度梯度图填充所有像素
 * @param grad_img 8UC1灰度图.
 * @param markers 32SC1的特征蒙版
 * @return Mat 单通道的蒙版
 */
Mat WaterShed_(const Mat &grad_img, const Mat &markers) {
    assert(!grad_img.empty());
    assert(grad_img.channels() == 1 && markers.channels() == 1);
    assert(grad_img.rows == markers.rows && grad_img.cols == markers.cols);
    const int graylevelnum = 256;
    Mat gray;
    grad_img.convertTo(gray, CV_8U);
    Mat res;
    markers.convertTo(res, CV_32S);
    vector<queue<point2D>> queues;
    queues.reserve(graylevelnum);
    //初始化后方便后面下标访问
    for (int i = 0; i < graylevelnum; i++) {
        queues.emplace_back(queue<point2D>());
    }
    int Rows = res.rows - 1, Cols = res.cols - 1;
    //首末行置SHED
    int *res_p_cur = res.ptr<int>(0);
    for (int i = 0; i < res.cols; i++) { res_p_cur[i] = SHED; }
    res_p_cur = res.ptr<int>(Rows);
    for (int i = 0; i < res.cols; i++) { res_p_cur[i] = SHED; }
    //遍历检测所有种子点并压入队列
    for (int i = 1; i < Rows; i++) {
        int *res_p_pre = res.ptr<int>(i - 1);
        res_p_cur = res.ptr<int>(i);//第i行
        int *res_p_next = res.ptr<int>(i + 1);
        auto *gray_p = gray.ptr<uchar>(i);
        //首末列置SHED
        res_p_cur[0] = SHED;
        res_p_cur[Cols] = SHED;

        for (int j = 1; j < Cols; j++) {
            if (res_p_cur[j] < 0) { res_p_cur[j] = 0; }
            int intensity = graylevelnum;
            if (res_p_cur[j] == 0 && (res_p_pre[j] > 0 || res_p_next[j] > 0 || res_p_cur[j - 1] > 0 ||
                                      res_p_cur[j + 1] > 0)) { intensity = gray_p[j]; }
            if (intensity < graylevelnum) {
                //压入点自身灰度值为零，周围存在种子点，根据灰度图上该点的灰度值区分不同序列
                queues[intensity].push({i, j});
                res_p_cur[j] = INQUEUE;
            }
        }
    }

    int ind = 0;
//        for (; ind < graylevelnum; ind++) {
//            if (!queues[ind].empty()) { break; }
//        }//直接由下面得到ind第一个值

    if (ind < graylevelnum)/*防止上述循环至ind==graylevelnum*/ {
        int label, t;
        while (true) {
            //直到把当前灰度值序列元素遍历完才寻找下一个不为空的灰度值序列
            if (queues[ind].empty()) {
                for (; ind < graylevelnum; ind++) {
                    if (!queues[ind].empty()) { break; }
                }
            }
            if (ind >= graylevelnum) { break; }
            //赋值后弹出
            point2D cur = queues[ind].front();
            queues[ind].pop();

            //倘若上下左右点存在大于零时的标记值不同，标记为边界，否则注水
            label = 0;
            t = res.ptr<int>(cur.r - 1)[cur.c];
            if (t > 0) label = t;
            t = res.ptr<int>(cur.r + 1)[cur.c];
            if (t > 0) {
                if (label == 0) { label = t; }
                else if (label != t) { label = SHED; }
            }
            t = res.ptr<int>(cur.r)[cur.c - 1];
            if (t > 0) {
                if (label == 0) { label = t; }
                else if (label != t) { label = SHED; }
            }
            t = res.ptr<int>(cur.r)[cur.c + 1];
            if (t > 0) {
                if (label == 0) { label = t; }
                else if (label != t) { label = SHED; }
            }

            assert(label != 0);//理论上压入队列内的点周围一定存在种子点，故在此断言
            res.ptr<int>(cur.r)[cur.c] = label;

            if (label != SHED) {
                //检查周围未标注的点，按灰度值压入队列中
                int r_ind = cur.r - 1, c_ind = cur.c;
                if (res.ptr<int>(r_ind)[c_ind] == 0) {
                    t = gray.ptr<uchar>(r_ind)[c_ind];
                    queues[t].push({r_ind, c_ind});
                    res.ptr<int>(r_ind)[c_ind] = INQUEUE;
                    ind = ind < t ? ind : t;//转向遍历新增的小灰度值队列
                }
                r_ind = cur.r + 1;
                c_ind = cur.c;
                if (res.ptr<int>(r_ind)[c_ind] == 0) {
                    t = gray.ptr<uchar>(r_ind)[c_ind];
                    queues[t].push({r_ind, c_ind});
                    res.ptr<int>(r_ind)[c_ind] = INQUEUE;
                    ind = ind < t ? ind : t;
                }
                r_ind = cur.r;
                c_ind = cur.c - 1;
                if (res.ptr<int>(r_ind)[c_ind] == 0) {
                    t = gray.ptr<uchar>(r_ind)[c_ind];
                    queues[t].push({r_ind, c_ind});
                    res.ptr<int>(r_ind)[c_ind] = INQUEUE;
                    ind = ind < t ? ind : t;
                }
                r_ind = cur.r;
                c_ind = cur.c + 1;
                if (res.ptr<int>(r_ind)[c_ind] == 0) {
                    t = gray.ptr<uchar>(r_ind)[c_ind];
                    queues[t].push({r_ind, c_ind});
                    res.ptr<int>(r_ind)[c_ind] = INQUEUE;
                    ind = ind < t ? ind : t;
                }
            }
        }
    }
    return res;
}

/**
 * @brief 根据特征矩阵填充彩色蒙版
 * @note C++11推荐随机数生成库 default_random_engine
 * @return
 */
void WaterShed::colorMask() {
    assert(resMask.channels() == 1 && basins > 0);//断言
    int min = 0, max = 255;//定义上下边界
    default_random_engine e(time(nullptr));//创建引擎
    uniform_int_distribution<int> u(min, max);//创建取值范围
    vector<RGB*> colors;
    colors.reserve(basins);
    RGB *RGB_p = nullptr;
    for (int i = 0; i < basins; ++i) {
        RGB_p = new RGB(u(e), u(e), u(e));
        colors.push_back(RGB_p);
    }
    for (int i = 1; i < result.rows - 1; ++i) {
        int *rM_p = resMask.ptr<int>(i);
        for (int j = 1; j < result.cols - 1; ++j) {
            if (rM_p[j] > 0) {
                result.ptr<Vec3b>(i)[j][0] = colors[rM_p[j] - 1]->b;
                result.ptr<Vec3b>(i)[j][1] = colors[rM_p[j] - 1]->g;
                result.ptr<Vec3b>(i)[j][2] = colors[rM_p[j] - 1]->r;
            }
        }
    }
    delete RGB_p;
    RGB_p = nullptr;
}

/**
 * @brief 根据特征矩阵填充彩色蒙版
 * @note 与上一种方法速度差别不大，但C++11推荐使用上面新的随机数生成方法
 * @return
 */
void WaterShed::ColorMask() {
    const int graylevelnum = 256;
    vector<RGB> colors;
    colors.reserve(basins);
    auto upper_bound = static_cast<double>(graylevelnum - 1);
    double step = upper_bound / ceil(pow(basins, 1 / 3.0));
    double semi_step = step / 2;
    double semi_semi_step = semi_step / 2;
    srand(256);
    int i = 0;
    for (double r = semi_step; r < upper_bound; r += step) {
        for (double g = semi_step; g < upper_bound; g += step) {
            for (double b = semi_step; b < upper_bound; b += step) {
                uchar R = r + (rand() * semi_step / RAND_MAX - semi_semi_step);
                uchar G = g + (rand() * semi_step / RAND_MAX - semi_semi_step);
                uchar B = b + (rand() * semi_step / RAND_MAX - semi_semi_step);
                colors.emplace_back(R, G, B);
                i++;
                if (i ==basins) { break; }
            }
            if (i ==basins) { break; }
        }
        if (i == basins) { break; }
    }

    for (i = 1; i < result.rows - 1; i++) {
        auto *p = result.ptr<Vec3b>(i);
        int *rM_p =resMask.ptr<int>(i);
        for (int j = 1; j <result.cols - 1; j++) {
            if (rM_p[j] > 0) {
                RGB cur = colors[rM_p[j] - 1];
                p[j][0] = cur.b;
                p[j][1] = cur.g;
                p[j][2] = cur.r;
            }
        }
    }
}

/**
 * @brief 主程序，手动给定种子分割
 * @return
 */
void WaterShed::waterShedWithMarkers() {
    this->markers();//绘制种子点

    this->resMask = WaterShed_(preProcess(), markersImg);
    this->colorMask();
    imshow("res", result);
    waitKey(0);
}

/**
 * @brief 自动寻找最小值
 * @return
 */
void WaterShed::findLocalMinimum() {
    Mat gray;
    cvtColor(src, gray, COLOR_BGR2GRAY);
    markersImg = Mat::zeros(gray.rows, gray.cols, CV_32S);
    //markersImg边缘标记为SHED
    int Rows = gray.rows - 1, Cols = gray.cols - 1;
    int *p = markersImg.ptr<int>(0);
    for (int i = 0; i < markersImg.cols; i++) { p[i] = SHED; }
    p = markersImg.ptr<int>(Rows);
    for (int i = 0; i < markersImg.cols; i++) { p[i] = SHED; }
    for (int i = 1; i < Rows; i++) {
        p = markersImg.ptr<int>(i);
        p[0] = SHED;
        p[Cols] = SHED;
    }
    queue<point2D> queue;//存储最小值点的位置信息(行，列)
    for (int i = 1; i < Rows; i++) {
        for (int j = 1; j < Cols; j++) {
            uchar val = gray.ptr<uchar>(i)[j];
            if (markersImg.ptr<int>(i)[j] == 0 && val <= gray.ptr<uchar>(i - 1)[j] &&
                val <= gray.ptr<uchar>(i + 1)[j] && val <= gray.ptr<uchar>(i)[j - 1] &&
                val <= gray.ptr<uchar>(i)[j + 1]) /*寻找灰度图中的局部最小值*/{
                basins++;
                markersImg.ptr<int>(i)[j] = basins;
                queue.push({i, j});
                //遍历最小值点周围上下左右四点，判断是否也为最小值,循环直至碰到灰度边界
                while (!queue.empty()) {
                    point2D cur = queue.front();
                    queue.pop();//赋值给cur后弹出
                    int R = cur.r - 1, C = cur.c;
                    int neighbor_mark = markersImg.ptr<int>(R)[C];
                    if (neighbor_mark == 0) {
                        if (val == gray.ptr<uchar>(R)[C]) {
                            markersImg.ptr<int>(R)[C] = basins;
                            queue.push({R, C});
                        }
                    }
                    R = cur.r + 1, C = cur.c;
                    neighbor_mark = markersImg.ptr<int>(R)[C];
                    if (neighbor_mark == 0) {
                        if (val == gray.ptr<uchar>(R)[C]) {
                            markersImg.ptr<int>(R)[C] = basins;
                            queue.push({R, C});
                        }
                    }

                    R = cur.r, C = cur.c - 1;
                    neighbor_mark = markersImg.ptr<int>(R)[C];
                    if (neighbor_mark == 0) {
                        if (val == gray.ptr<uchar>(R)[C]) {
                            markersImg.ptr<int>(R)[C] = basins;
                            queue.push({R, C});
                        }
                    }

                    R = cur.r, C = cur.c + 1;
                    neighbor_mark = markersImg.ptr<int>(R)[C];
                    if (neighbor_mark == 0) {
                        if (val == gray.ptr<uchar>(R)[C]) {
                            markersImg.ptr<int>(R)[C] = basins;
                            queue.push({R, C});
                        }
                    }
                }
            }
        }
    }
}

/**
 * @brief 自动寻找最小值进行图像分割(主程序)
 * @note 易出现过分割现象,仅图一乐，没多大实际意义
 * @return
 */
void WaterShed::waterShedAuto() {
    findLocalMinimum();
//        cout << basins << endl;
    this->resMask = WaterShed_(preProcess(), markersImg);
    this->colorMask();
    imshow("res", result);
    waitKey(0);
}

/**
 * @brief 自动对视频进行图像分割(主程序)
 * @note 易出现过分割现象,视频懒得测试
 * @return
 */
void WaterShed::waterShedVideo(const string &video_name) {
    VideoCapture cap;
    cap.open(video_name);
    //等比例保存
    Size size = Size(cap.get(CAP_PROP_FRAME_WIDTH), cap.get(CAP_PROP_FRAME_HEIGHT));
    VideoWriter writer("./WaterShed/out.mp4", VideoWriter::fourcc('X', '2', '6', '4'), cap.get(CAP_PROP_FPS), size,
                       true);//H.264格式
    if (writer.isOpened()) {
        if (cap.isOpened()) {
            Mat frame;
            while (true) {
                cap >> frame;
                if (frame.empty()) {
                    break;
                }
                this->Init(frame);
                this->findLocalMinimum();
                this->resMask = WaterShed_(preProcess(), markersImg);
                this->colorMask();
                imshow("res", result);
                waitKey(1);
                writer << result;
            }
            cap.release();
            writer.release();
        }
    }
}